package com.alibaba.cloud.ai.controller;

import com.alibaba.cloud.ai.connector.config.DbConfig;
import com.alibaba.cloud.ai.dto.GraphRequest;
import com.alibaba.cloud.ai.dto.InitSchemaDto;
import com.alibaba.cloud.ai.entity.ChatMessage;
import com.alibaba.cloud.ai.entity.Datasource;
import com.alibaba.cloud.ai.graph.CompileConfig;
import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.request.SchemaInitRequest;
import com.alibaba.cloud.ai.service.AgentDatasourceService;
import com.alibaba.cloud.ai.service.ChatMessageService;
import com.alibaba.cloud.ai.service.ChatSessionService;
import com.alibaba.cloud.ai.service.datasource.DatasourceService;
import com.alibaba.cloud.ai.service.graph.GraphService;
import com.alibaba.cloud.ai.service.vectorstore.AgentVectorStoreService;
import com.alibaba.cloud.ai.util.SchemaProcessorUtil;
import com.alibaba.cloud.ai.vo.ApiResponse;
import com.alibaba.cloud.ai.vo.GraphNodeResponse;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;

import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Optional;

import static com.alibaba.cloud.ai.constant.Constant.*;

/**
 * @author taoshujian
 * @version 1.0
 * @since 2025-11-07 13:53:31
 */
@Slf4j
@RestController
@RequestMapping("/tsj/api")
public class TsjController {

    @Resource
    private AgentDatasourceService agentDatasourceService;
    @Resource
    private ChatMessageService chatMessageService;
    @Resource
    private ChatSessionService chatSessionService;
    @Resource
    private GraphService graphService;
    @Autowired
    private AgentVectorStoreService simpleVectorStoreService;
    @Autowired
    private DatasourceService datasourceService;

    private final CompiledGraph compiledGraph;

    @Autowired
    @Qualifier("nl2sqlGraph2")
    private StateGraph stateGraph2;


    @Autowired
    public TsjController(@Qualifier("nl2sqlGraph") StateGraph stateGraph) throws GraphStateException {
        this.compiledGraph = stateGraph.compile();
        //this.compiledGraph = stateGraph.compile(CompileConfig.builder().interruptBefore(HUMAN_FEEDBACK_NODE).build());
        this.compiledGraph.setMaxIterations(100);
    }

    @GetMapping("/nl2sql")
    public String nl2sql(@RequestParam String query) throws Exception {
        //return graphService.nl2sql(query,2+"");
        initSchema(3);
        Map<String, Object> inputs = Map.of(IS_ONLY_NL2SQL, true, INPUT_KEY, query, AGENT_ID, 2+"");
        Optional<OverAllState> call = compiledGraph.call(inputs);
        OverAllState state = call.orElseThrow();
        return state.value(ONLY_NL2SQL_OUTPUT).toString();
    }

    /**
     * http://127.0.0.1:9102/tsj/api/search?query=近一年商品销量排行榜
     * @param query
     * @return
     * @throws Exception
     */
    @GetMapping("/search")
    public String search(@RequestParam String query) throws Exception {

//        SchemaInitRequest schemaInitRequest = new SchemaInitRequest();

//        Datasource datasource = datasourceService.getDatasourceById(3);
//        // Create database configuration
//        DbConfig dbConfig = SchemaProcessorUtil.createDbConfigFromDatasource(datasource);
//        schemaInitRequest.setDbConfig(dbConfig);
//        schemaInitRequest
//                .setTables(Arrays.asList("categories", "order_items", "orders", "products", "users", "product_categories"));
//        simpleVectorStoreService.schema(2+"",schemaInitRequest);

        CompiledGraph compiledGraph2 = stateGraph2.compile();
        compiledGraph2.setMaxIterations(100);

        Map<String, Object> inputs = Map.of(INPUT_KEY, query, AGENT_ID, "2");
        Optional<OverAllState> call = compiledGraph.call(inputs);
        OverAllState state = call.orElseThrow();
        return state.value(RESULT).toString();
    }

    @GetMapping("/search2")
    public String search2(@RequestParam String query) throws Exception {
        initSchema(3);
        CompiledGraph compiledGraph2 = stateGraph2.compile();
        compiledGraph2.setMaxIterations(100);

        Map<String, Object> inputs = Map.of(INPUT_KEY, query, AGENT_ID, "2");
        Optional<OverAllState> call = compiledGraph2.call(inputs);
        OverAllState state = call.orElseThrow();
        return state.value(RESULT).toString();
    }



    /**
     * 初始化数据源
     */
    @PostMapping("/init")
    public ResponseEntity<ApiResponse> initSchema(Integer datasourceId) {
        try {

            if (datasourceId == null) {
                return ResponseEntity.badRequest().body(ApiResponse.error("数据源ID不能为空"));
            }

            List<String> tables = agentDatasourceService.getDatasourceTables(datasourceId);

            if (tables == null || tables.isEmpty()) {
                return ResponseEntity.badRequest().body(ApiResponse.error("表列表不能为空"));
            }

            // Execute Schema initialization
            Boolean result = agentDatasourceService.initializeSchemaForAgentWithDatasource(2L, datasourceId,
                    tables);

            if (result) {
                log.info("Successfully initialized schema for agent: {}, tables: {}", 2, tables.size());
                return ResponseEntity.ok(ApiResponse.success("Schema初始化成功"));
            }
            else {
                return ResponseEntity.internalServerError().body(ApiResponse.error("Schema初始化失败"));
            }
        }
        catch (Exception e) {
            log.error("Failed to initialize schema for agent: {}", 2L, e);
            return ResponseEntity.internalServerError().body(ApiResponse.error("Schema初始化失败：" + e.getMessage()));
        }
    }

    /**
     * 新增会话消息
     */
    @PostMapping("/saveMessage")
    public ResponseEntity<ChatMessage> saveMessage(@RequestBody ChatMessage message) {
        try {
            message.setSessionId("ef80d54a-c894-4208-af65-99928a272f3f");
            message.setRole("user");
            message.setMessageType("text");
            ChatMessage savedMessage = chatMessageService.saveMessage(message);

            chatSessionService.updateSessionTime(message.getSessionId());

            return ResponseEntity.ok(savedMessage);
        }
        catch (Exception e) {
            log.error("Save message error for session {}: {}", message.getSessionId(), e.getMessage(), e);
            return ResponseEntity.internalServerError().build();
        }
    }

    @GetMapping(value = "/stream/search", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<GraphNodeResponse>> streamSearch(String query,HttpServletResponse response) {
        // Set SSE-related HTTP headers
        response.setCharacterEncoding("UTF-8");
        response.setContentType("text/event-stream");
        response.setHeader("Cache-Control", "no-cache");
        response.setHeader("Connection", "keep-alive");
        response.setHeader("Access-Control-Allow-Origin", "*");
        response.setHeader("Access-Control-Allow-Headers", "Cache-Control");

        //Project Reactor 中的 Sinks 是 Spring WebFlux 中用于响应式编程的一个组件，特别是用于 Server-Sent Events (SSE) 的场景。
        Sinks.Many<ServerSentEvent<GraphNodeResponse>> sink = Sinks.many().unicast().onBackpressureBuffer();

        GraphRequest request = GraphRequest.builder()
                .agentId(2+"")
                .threadId(null)
                .query(query)
                .humanFeedback(false)
                .humanFeedbackContent(null)
                .rejectedPlan(false)
                .nl2sqlOnly(false)
                .plainReport(true)
                .build();
        graphService.graphStreamProcess(sink, request);

        return sink.asFlux()
                .doOnSubscribe(subscription -> log.info("Client subscribed to stream"))
                .doOnCancel(() -> log.info("Client disconnected from stream"))
                .doOnError(e -> log.error("Error occurred during streaming: ", e))
                .doOnComplete(() -> log.info("Stream completed successfully"));
    }

    /**
     * 查询会话消息
     * @param sessionId
     * @return
     */
    @PostMapping("/getMessage")
    public ResponseEntity<List<ChatMessage>> getMessage(String sessionId) {
        List<ChatMessage> messages = chatMessageService.findBySessionId(sessionId);
        return ResponseEntity.ok(messages);
    }
}
